import math
import cv2
import torch
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
import torch.nn.functional as F

from datasets.table.dataloader import LoadTableImageAndLabels
from models.assembly.segmentation_table import Segmentation_Model


def save_tensor(tensor, i, e):
    np_array = tensor[0].detach().cpu().numpy().transpose(1, 2, 0)
    #     mx, mn = np_array.max(), np_array.min()
    #     arr = (np_array - mn) / (mx - mn) * 255
    cv2.imwrite(str(e) + '-' + str(i) + '.jpg', np_array * 255)


def main():
    data_fd = r'F:\datasets\ocr\internal\mask-label'
    # data_fd = '/mnt/data/xp/datasets/table/images'
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    dataset = LoadTableImageAndLabels(data_fd)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, num_workers=0, shuffle=True, pin_memory=True)
    model = Segmentation_Model().to(device)
    epochs = 100

    optimizer = optim.Adam(model.parameters(), lr=0.001)  # , weight_decay=5e-4)

    lf = lambda x: (((1 + math.cos(x * math.pi / epochs)) / 2) ** 1.0) * 0.9 + 0.1  # cosine
    scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
    scheduler.last_epoch = epochs - 1  # do not move
    for epoch in range(epochs):
        model.train()
        print('epoch: ', epoch)
        for i, (imgs, targets) in enumerate(dataloader):
            imgs = imgs.to(device).float() / 255.0
            pred = model(imgs)
            loss = F.mse_loss(pred, targets.to(device))
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            if i % 10 == 0:
                save_tensor(pred, i, epoch)
                print(i, pred.max(), pred.min(), loss)
        scheduler.step()
        torch.save(model.state_dict(), str(epoch) + '.pth')


if __name__ == '__main__':
    main()
